// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>

namespace kumo::pollux::functions::sparksql::test {
namespace {

class UuidTest : public SparkFunctionBaseTest {
 protected:
  std::optional<std::string> uuidOne(int64_t seed, int32_t partitionIndex) {
    setSparkPartitionId(partitionIndex);
    return evaluateOnce<std::string>(
        fmt::format("uuid({})", seed), make_row_vector(ROW({}), 1));
  }

  VectorPtr uuidMany(int64_t seed, int32_t partitionIndex, int32_t batchSize) {
    setSparkPartitionId(partitionIndex);
    auto exprSet = compileExpression(fmt::format("uuid({})", seed), ROW({}));
    return evaluate(*exprSet, make_row_vector(ROW({}), batchSize));
  }
};

TEST_F(UuidTest, basic) {
  // With same partitionIndex, same seed always produces same result.
  EXPECT_EQ(uuidOne(100, 1), uuidOne(100, 1));
  EXPECT_EQ(uuidOne(2, 1234), uuidOne(2, 1234));

  EXPECT_NE(uuidOne(9, 1), uuidOne(100, 1));
  EXPECT_NE(uuidOne(9, 1), uuidOne(9, 2));
  EXPECT_NE(uuidOne(100, 1), uuidOne(99, 20));

  // Expected result comes from Spark.
  EXPECT_EQ(uuidOne(0, 0), std::string("8c7f0aac-97c4-4a2f-b716-a675d821ccc0"));

  pollux::test::assertEqualVectors(
      uuidMany(123, 1233, 100), uuidMany(123, 1233, 100));
  pollux::test::assertEqualVectors(
      uuidMany(321, 1233, 33), uuidMany(321, 1233, 33));
}

TEST_F(UuidTest, nonConstantSeed) {
  setSparkPartitionId(0);

  POLLUX_ASSERT_THROW(
      evaluateOnce<std::string>("uuid(c0)", std::optional<int64_t>(123)),
      "seed argument must be constant");
}

} // namespace
} // namespace kumo::pollux::functions::sparksql::test
